﻿using System;
using System.Collections.Concurrent;
using System.Collections.Generic;
using System.Linq;
using System.Threading.Tasks;
using System.Windows.Forms;
using Ast.Core.Forms;
using Ast.Core.Services;
using Ast.Core.Services.Extensions;
using Ast.Core.Services.ObjectDefinitions;
using Ast.Framework.Forms;
using Castle.Core.Logging;

namespace Ast
{
    [FormMetadata(FormType.Routine, IsTabForm = true)]
    public partial class RoutineForm : Form
    {
        private readonly IPsLoadedModuleListService _psLoadedModuleListService;
        private readonly IKernelRoutineService _kernelRoutineService;
        private readonly IDeviceController _deviceController;

        public RoutineForm(IPsLoadedModuleListService psLoadedModuleListService, IKernelRoutineService kernelRoutineService, IDeviceController deviceController)
        {
            _psLoadedModuleListService = psLoadedModuleListService;
            _kernelRoutineService = kernelRoutineService;
            _deviceController = deviceController;

            InitializeComponent();

            Logger = NullLogger.Instance;
        }

        public ILogger Logger { get; set; }

        private void RoutineForm_Load(object sender, EventArgs e)
        {
            // list view sorter
            lvRoutines.ApplyColumnSorter();

            // load
            btnReload.PerformClick();
        }

        private async void btnReload_Click(object sender, EventArgs e)
        {
            // load data
            var kernelModules = _psLoadedModuleListService.GetAllModules();
            var items = new ConcurrentBag<ListViewItem>();
            await Task.WhenAll(
                Task.Run(() => _kernelRoutineService.ListProcessNotifyRoutines(routine =>
                {
                    AddViewItem(kernelModules, routine.Function, "CreateProcess", items);
                    return true;
                })),
                Task.Run(() => _kernelRoutineService.ListThreadNotifyRoutines(routine =>
                {
                    AddViewItem(kernelModules, routine.Function, "CreateThread", items);
                    return true;
                })),
                Task.Run(() => _kernelRoutineService.ListImageNotifyRoutines(routine =>
                {
                    AddViewItem(kernelModules, routine.Function, "LoadImage", items);
                    return true;
                })),
                Task.Run(() => _kernelRoutineService.ListRegisterCallbacks(item =>
                {
                    AddViewItem(kernelModules, item.Function, "CmpCallback", items);
                    return true;
                })),
                Task.Run(() => _kernelRoutineService.ListShutdownNotifications(item =>
                {
                    var deviceObject = _deviceController.ReadAs<DeviceObject>(item.DeviceObject);
                    var driverObject = _deviceController.ReadAs<DriverObject>(deviceObject.DriverObject);
                    AddViewItem(kernelModules, driverObject.MajorFunction[0x10], "Shutdown", items); // #define IRP_MJ_SHUTDOWN                 0x10
                    return true;
                })),
                Task.Run(() => _kernelRoutineService.ListBugCheckCallbacks(item =>
                {
                    AddViewItem(kernelModules, item.CallbackRoutine, "BugCheckCallback", items);
                    return true;
                })),
                Task.Run(() => _kernelRoutineService.ListBugCheckReasonCallbacks(item =>
                {
                    AddViewItem(kernelModules, item.CallbackRoutine, "BugCheckReasonCallback", items);
                    return true;
                })),
                Task.Run(() => _kernelRoutineService.ListFsRegistrationChangeRoutines(item =>
                {
                    AddViewItem(kernelModules, item.DriverNotificationRoutine, "FsRegistrationChange", items);
                    return true;
                })),
                Task.Run(() => _kernelRoutineService.ListPowerSettingCallbacks(item =>
                {
                    AddViewItem(kernelModules, item.Callback, "PowerSettingCallback", items);
                    return true;
                })),
                Task.Run(() => _kernelRoutineService.ListPlugPlayNotifications(item =>
                {
                    AddViewItem(kernelModules, item.CallbackRoutine, "PlugPlay", items);
                    return true;
                })),
                Task.Run(() => _kernelRoutineService.ListLogonSessionTerminatedRoutines(item =>
                {
                    AddViewItem(kernelModules, item.CallbackRoutine, "SeFileSystem", items);
                    return true;
                })),
                Task.Run(() =>
                {
                    var legoNotifyRoutine = _kernelRoutineService.GetLegoNotifyRoutine();
                    if (legoNotifyRoutine != IntPtr.Zero)
                    {
                        AddViewItem(kernelModules, legoNotifyRoutine, "Lego", items);
                    }
                })
                ).ConfigureAwait(false);

            Logger.InfoFormat("found {0} routines", items.Count);

            // update view
            Invoke(new Action(() =>
            {
                lvRoutines.BeginUpdate();
                lvRoutines.Items.Clear();
                lvRoutines.Items.AddRange(items.ToArray());
                lvRoutines.AutoResizeColumns();
                lvRoutines.EndUpdate();
            }));
        }

        private void AddViewItem(IEnumerable<KernelModule> kernelModules, IntPtr function, string type, ConcurrentBag<ListViewItem> itemContainer)
        {
            // item tag
            var item = new RoutineViewItem
            {
                Function = function,
                Type = type,
            };
            var funcAdr = (uint) function.ToInt32();
            item.Module = kernelModules.FirstOrDefault(x => funcAdr >= x.StartAddress && funcAdr <= x.EndAddress);

            // view item
            var listViewItem = new ListViewItem(new[]
            {
                "0x" + item.Function.ToString("x8"),
                item.Type,
                item.Module != null ? item.Module.Path : string.Empty
            })
            {
                Tag = item
            };

            itemContainer.Add(listViewItem);
        }
    }

    public class RoutineViewItem
    {
        public IntPtr Function { get; set; }
        public string Type { get; set; }
        public KernelModule Module { get; set; }
    }
}